package com.LachesismStorm.lachesism.ai.kimi;


import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.http.ContentType;
import cn.hutool.http.Header;
import cn.hutool.http.HttpRequest;
import cn.hutool.http.Method;
import cn.hutool.json.JSON;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import com.google.gson.JsonObject;
import lombok.NonNull;
import lombok.SneakyThrows;
import okhttp3.*;

import java.io.BufferedReader;
import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.TimeUnit;

/**
 * @Author lachesism
 * @Date 2024-08-01 21:12
 **/
public class MoonshotAiUtils {

    private static final String API_KEY = "sk-RSrKcdXsSfCT6W12bHrRh5OKWcTqntLgdFnjkScSRSM0Q6vg";
    private static final String MODELS_URL = "https://api.moonshot.cn/v1/models";
    private static final String FILES_URL = "https://api.moonshot.cn/v1/files";
    private static final String ESTIMATE_TOKEN_COUNT_URL = "https://api.moonshot.cn/v1/tokenizers/estimate-token-count";
    private static final String CHAT_COMPLETION_URL = "https://api.moonshot.cn/v1/chat/completions";

    public static String getModelList() {
        return getCommonRequest(MODELS_URL)
                .execute()
                .body();
    }

    public static String uploadFile(@NonNull File file) {
        return getCommonRequest(FILES_URL)
                .method(Method.POST)
                .header("purpose", "file-extract")
                .form("file", file)
                .execute()
                .body();
    }

    public static String getFileList() {
        return getCommonRequest(FILES_URL)
                .execute()
                .body();
    }

    public static String deleteFile(@NonNull String fileId) {
        return getCommonRequest(FILES_URL + "/" + fileId)
                .method(Method.DELETE)
                .execute()
                .body();
    }

    public static String getFileDetail(@NonNull String fileId) {
        return getCommonRequest(FILES_URL + "/" + fileId)
                .execute()
                .body();
    }

    public static String getFileContent(@NonNull String fileId) {
        return getCommonRequest(FILES_URL + "/" + fileId + "/content")
                .execute()
                .body();
    }

    public static String estimateTokenCount(@NonNull String model, @NonNull List<Message> messages) {
        String requestBody = new JSONObject()
                .putOpt("model", model)
                .putOpt("messages", messages)
                .toString();
        return getCommonRequest(ESTIMATE_TOKEN_COUNT_URL)
                .method(Method.POST)
                .header(Header.CONTENT_TYPE, ContentType.JSON.getValue())
                .body(requestBody)
                .execute()
                .body();
    }

    public static void main(String[] args) {
        List<Message> messages = new ArrayList<>();
        Message message = new Message();
        message.setRole("user");
        message.setContent("我上一句话说的是什么");
        messages.add(message);
        System.out.println(chat("moonshot-v1-32k", messages));
    }

    @SneakyThrows
    public static List<String> chat(@NonNull String model, @NonNull List<Message> messages) {
        System.out.println('\n' + "kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi");
        System.out.println(messages);
        System.out.println("kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi" + '\n');
        List<String> responseMessages = new ArrayList<>();
        String requestBody = new JSONObject()
                .putOpt("model", model)
                .putOpt("messages", messages)
                .putOpt("stream", false)
                .toString();
        Request okhttpRequest = new Request.Builder()
                .url(CHAT_COMPLETION_URL)
                .post(RequestBody.create(requestBody, MediaType.get(ContentType.JSON.getValue())))
                .addHeader("Authorization", "Bearer " + API_KEY)
                .build();

        OkHttpClient client = new OkHttpClient.Builder()
                .connectTimeout(120, TimeUnit.SECONDS)
                .readTimeout(120, TimeUnit.SECONDS)
                .build();
        Call call = client.newCall(okhttpRequest);
        Response okhttpResponse = call.execute();
        if (!okhttpResponse.isSuccessful()) {
            throw new IOException("Unexpected code " + okhttpResponse);
        }
        JSONObject jsonResponse = (JSONObject) JSONUtil.parse(okhttpResponse.body().string());
        System.out.println('\n' + "kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi");
        System.out.println(jsonResponse.getJSONArray("choices"));
        System.out.println("kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi=kimi" + '\n');
        jsonResponse.getJSONArray("choices").forEach(choice -> {
            JSONObject messageObj = ((JSONObject) choice).getJSONObject("message");
            String content = (String) messageObj.get("content");
            responseMessages.add(content);
        });
        return responseMessages;
    }


    private static HttpRequest getCommonRequest(@NonNull String url) {
        return HttpRequest.of(url).header(Header.AUTHORIZATION, "Bearer " + API_KEY);
    }

}